package parser;

import java.util.ArrayList;
import java.util.List;
import static parser.MySimpleNode.jjtSetChild;
import static parser.MySimpleNode.jjtAppendChild;
import static parser.MySimpleNode.jjtGetChildNum;
import static parser.MySimpleNode.errorExit;

public class JoinHelp implements XQueryParserTreeConstants
{

	public static void changeFLWRtoJoin(ASTXQueryFLWR inFLWR)
	{
		ASTFor inFor = (ASTFor) inFLWR.jjtGetChild(0);
		ASTWhere inWhere = (ASTWhere) inFLWR.jjtGetChild(2);
		List<List<String>> magicLists = makeMagicLists(inFor);

		if (magicLists.size() < 2)
			errorExit("You need at least two lists to join, you fool!");

		// Build basic structure for outer FLWR
		ASTXQueryFLWR retval = new ASTXQueryFLWR(JJTXQUERYFLWR);

		// Will setup the for later in this function
		ASTFor outerFor = new ASTFor(JJTFOR);
		outerFor.vars.add("$tuple");

		// The let and Where are just empty
		ASTLet outerLet = new ASTLet(JJTLET);
		ASTWhere outerWhere = new ASTWhere(JJTWHERE);

		// The new return statement is just the old one (after processing)
		ASTReturn outerReturn = (ASTReturn) inFLWR.jjtGetChild(3);

		// Setup parent/child relations for outer FLWR
		jjtSetChild(retval, outerFor, 0);
		jjtSetChild(retval, outerLet, 1);
		jjtSetChild(retval, outerWhere, 2);
		jjtSetChild(retval, outerReturn, 3);

		ASTXQueryJoin joinXQ = setupNestedJoins(inFLWR, inWhere, magicLists,
				null);
		jjtSetChild(outerFor, joinXQ, 0);

		// Setup the outerReturn
		makeOuterJoinReturn(outerReturn);

		// Replace the old FLWR with the new one using join
		int childNum = jjtGetChildNum(inFLWR.jjtGetParent(), inFLWR);
		jjtSetChild(inFLWR.jjtGetParent(), retval, childNum);

	}

	private static ASTXQueryJoin setupNestedJoins(ASTXQueryFLWR inFLWR,
			ASTWhere inWhere, List<List<String>> magicLists,
			ASTXQueryJoin buildJoin)
	{

		// Done building the nested joins.
		if (magicLists.size() == 1)
			return buildJoin;

		List<String> relVars1 = magicLists.get(0);
		List<String> relVars2 = magicLists.get(1);
		ASTXQuery R1XQ = null;
		ASTXQueryFLWR R2XQ = null;
		// Returns the two lists of join variables
		List<String>[] joinLists = findJoinVars(relVars1, relVars2, inWhere);

		// Innermost join is processed differently
		if (buildJoin == null)
			R1XQ = makeInnerJoinFor(inFLWR, relVars1);
		else
			R1XQ = buildJoin;

		R2XQ = makeInnerJoinFor(inFLWR, relVars2);
		ASTXQueryJoin joinXQ = createJoin(R1XQ, R2XQ, joinLists[0],
				joinLists[1]);

		// Combine the first two magic lists.
		magicLists.get(0).addAll(magicLists.get(1));
		magicLists.remove(1);
		return setupNestedJoins(inFLWR, inWhere, magicLists, joinXQ);
	}

	private static ASTXQueryJoin createJoin(ASTXQuery xquery1,
			ASTXQueryFLWR flwr2, List<String> joinTags1, List<String> joinTags2)
	{
		ASTXQueryJoin retval = new ASTXQueryJoin(JJTXQUERYJOIN);
		retval.tags1 = joinTags1;
		retval.tags2 = joinTags2;
		jjtSetChild(retval, xquery1, 0);
		jjtSetChild(retval, flwr2, 1);
		return retval;
	}

	/**
	 * 
	 * @param variable
	 *            The variable that needs to be rewritten for the inner for loop
	 * 
	 * @return The tree created by replacing instances of $foo with
	 *         $tupple/foo/*
	 */
	private static ASTXQuerySlashRP fixJoinVar(ASTXQueryVar variable)
	{

		boolean slashRP = variable.jjtGetParent().id == JJTXQUERYSLASHRP;
		boolean dSlashRP = variable.jjtGetParent().id == JJTXQUERYDSLASHRP;

		// What was a variable is now a path from $tuple (with $ stripped off)
		String oldVar = variable.varName;
		oldVar = oldVar.substring(1);

		// XQuerySLashRP connecting $tuple to everything else
		ASTXQuerySlashRP retval = new ASTXQuerySlashRP(JJTXQUERYSLASHRP);

		// A new variable created to index tuples
		ASTXQueryVar tuppleVar = new ASTXQueryVar(JJTXQUERYVAR);
		tuppleVar.varName = "$tuple";

		// The slash connecting varTag to the rest of the path (wild or
		// otherwise)
		ASTRelativePathSlash varTagSlash = new ASTRelativePathSlash(
				JJTRELATIVEPATHSLASH);

		// Connect $tupple with rp/
		jjtSetChild(retval, tuppleVar, 0);
		jjtSetChild(retval, varTagSlash, 1);

		// What was a variable is now the start of a path from $tuple
		ASTRelativePathTag varTag = new ASTRelativePathTag(JJTRELATIVEPATHTAG);
		varTag.tagName = oldVar;
		jjtSetChild(varTagSlash, varTag, 0);

		// Needed to remove extra level of indirection created from tuple
		// creation
		ASTRelativePathWild wild = new ASTRelativePathWild(JJTRELATIVEPATHWILD);

		// If there's no other subtree, just make foo/*
		if (!slashRP && !dSlashRP)
			jjtSetChild(varTagSlash, wild, 1);

		// Otherwise we need to connect the old subtree
		else
		{
			ASTRelativePath subTreeConnector = null;

			// If slash, make the slash connecting the wild indirection to the
			// old subtree
			if (slashRP)
				subTreeConnector = new ASTRelativePathSlash(
						JJTRELATIVEPATHSLASH);
			// Otherwise make a dslash conncet the wild indirection to the old
			// subtree

			else
				subTreeConnector = new ASTRelativePathDSlash(
						JJTRELATIVEPATHDSLASH);

			jjtSetChild(varTagSlash, subTreeConnector, 1);
			jjtSetChild(subTreeConnector, wild, 0);
			// The subtree can be found as the other child of $foo's parent node
			jjtSetChild(subTreeConnector, variable.jjtGetParent()
					.jjtGetChild(1), 1);
		}

		return retval;
	}

	/**
	 * Takes in a list of relevant variable and an original for loop and where
	 * statement And writes the for loop that deals with the relevant variables
	 * inside a new join
	 * 
	 * @param inFor
	 * @param variable2
	 * @return
	 */
	private static ASTXQueryFLWR makeInnerJoinFor(ASTXQueryFLWR inFLWR,
			List<String> relVars)
	{

		ASTFor inFor = (ASTFor) inFLWR.jjtGetChild(0);
		ASTWhere inWhere = (ASTWhere) inFLWR.jjtGetChild(2);

		ASTXQueryFLWR retval = new ASTXQueryFLWR(JJTXQUERYFLWR);
		// For loop used in new join using vars in relVars
		ASTFor newFor = new ASTFor(JJTFOR);
		newFor.vars.addAll(relVars);
		// assuming no repeated var name in relVars,
		for (String var : relVars)
		{
			// Get the index of the current relevant variable in the original
			// forloop
			if (!inFor.vars.contains(var))
				errorExit("The variable " + var
						+ " is not in the original query");
			int varNum = inFor.vars.indexOf(var);
			ASTXQuery varQuery = (ASTXQuery) inFor.jjtGetChild(varNum);
			jjtAppendChild(newFor, varQuery);
		}
		jjtSetChild(retval, newFor, 0);

		// We are not dealing with FLWR sentences with LET statements
		ASTLet newLet = new ASTLet(JJTLET);
		jjtSetChild(retval, newLet, 1);

		ASTWhere newWhere = makeInnerJoinWhere(inWhere, relVars);
		jjtSetChild(retval, newWhere, 2);

		ASTReturn newReturn = makeInnerJoinReturn(relVars);
		jjtSetChild(retval, newReturn, 3);

		return retval;
	}

	/**
	 * Builds the tuple return statemnt
	 * <tuple>{<var>{}</var>...<var>{}</var>}</tuple> based on a relevant
	 * variable list
	 * 
	 * @param relVars
	 *            The variables to put into this form
	 * @return The return statement using varaibels of that form
	 */
	private static ASTReturn makeInnerJoinReturn(List<String> relVars)
	{
		if (relVars.isEmpty())
			errorExit("There should be at least one relevant "
					+ "variable in each partition to perform a join");
		// Set up the outer tag labeled tuple
		ASTReturn retval = new ASTReturn(JJTRETURN);
		ASTXQueryTag tupleTag = new ASTXQueryTag(JJTXQUERYTAG);
		tupleTag.tagName = "tuple";
		jjtSetChild(retval, tupleTag, 0);

		// If only one relVar, don't bother with commas
		if (relVars.size() == 1)
		{
			String var = relVars.get(0);
			ASTXQueryTag varTag = new ASTXQueryTag(JJTXQUERYTAG);
			varTag.tagName = var.substring(1);
			ASTXQueryVar varVar = new ASTXQueryVar(JJTXQUERYVAR);
			varVar.varName = var;

			jjtSetChild(varTag, varVar, 0);
			jjtSetChild(tupleTag, varTag, 0);
			return retval;
		}

		// If there is more than one relVar, process all but the last two the
		// same
		List<ASTXQueryComma> commaList = new ArrayList<ASTXQueryComma>();

		// There should be one less comma than there are relvars
		for (int i = 0; i < relVars.size() - 1; i++)
			commaList.add(new ASTXQueryComma(JJTXQUERYCOMMA));

		// Connect the first comma to the tupple tag
		jjtSetChild(tupleTag, commaList.get(0), 0);

		int index = 0;
		while (index < commaList.size() - 1)
		{
			ASTXQueryComma parentComma = commaList.get(index);
			ASTXQueryComma childComma = commaList.get(index + 1);
			String var = relVars.get(index);

			ASTXQueryTag varTag = makeVarTagSubtree(var);

			// Set up parent Comma's children nodes
			jjtSetChild(parentComma, varTag, 0);
			jjtSetChild(parentComma, childComma, 1);
			index++;
		}

		// The last comma node gets the last two tag nodes for children
		String penultVar = relVars.get(index);
		String finalVar = relVars.get(index + 1);

		ASTXQueryTag penultVarTag = makeVarTagSubtree(penultVar);
		ASTXQueryTag finalVarTag = makeVarTagSubtree(finalVar);
		jjtSetChild(commaList.get(index), penultVarTag, 0);
		jjtSetChild(commaList.get(index), finalVarTag, 1);

		return retval;
	}

	private static ASTWhere makeInnerJoinWhere(ASTWhere inWhere,
			List<String> relVars)
	{
		// If old where was empty, clearly new one is as well;
		if (inWhere.jjtGetNumChildren() == 0)
			return inWhere;

		ASTWhere retval = new ASTWhere(JJTWHERE);
		List<ASTCondEQ> eqList = new ArrayList<ASTCondEQ>();
		List<ASTCondAnd> andList = new ArrayList<ASTCondAnd>();

		List<ASTCondEQ> flattenedWhere = getEqFromWhere(inWhere,
				new ArrayList<ASTCondEQ>());

		eqList = (extractInnerWhereEQs(flattenedWhere, relVars));

		if (eqList.size() == 0)
			return retval;
		if (eqList.size() == 1)
		{
			jjtSetChild(retval, eqList.get(0), 0);
			return retval;
		}

		// Make AND node list that is one AND node short of the eqNodeList
		for (int i = 0; i < eqList.size() - 1; i++)
			andList.add(new ASTCondAnd(JJTCONDAND));

		int index = 0;
		// The last and is processed differently since its children are both eq
		// nodes
		while (index < andList.size() - 1)
		{
			ASTCondAnd parentAnd = andList.get(index);
			ASTCondAnd childAnd = andList.get(index + 1);
			ASTCondEQ childEQ = eqList.get(index);
			jjtSetChild(parentAnd, childEQ, 0);
			jjtSetChild(parentAnd, childAnd, 1);
			index++;
		}
		// The last AND node gets the last two EQ nodes for children
		jjtSetChild(andList.get(index), eqList.get(index), 0);
		jjtSetChild(andList.get(index), eqList.get(index + 1), 1);
		jjtSetChild(retval, andList.get(0), 0);
		return retval;
	}

	/**
	 * Take a varName $foo and makes a tag tree <foo>{$foo}</foo> Used for
	 * MakeInnerJoinReturn
	 * 
	 * @param varName
	 *            The variable to turn into the tree
	 * @return The tree created
	 */
	private static ASTXQueryTag makeVarTagSubtree(String varName)
	{
		ASTXQueryTag varTag = new ASTXQueryTag(JJTXQUERYTAG);
		varTag.tagName = varName.substring(1);
		ASTXQueryVar varVar = new ASTXQueryVar(JJTXQUERYVAR);
		varVar.varName = varName;
		jjtSetChild(varTag, varVar, 0);
		return varTag;
	}

	/**
	 * Take a flattened where tree, and prune off the parts that aren't selects
	 * on the relevant variables, or joins between relevant variables.
	 * 
	 * Used for MakeInnerJoinWhere
	 * 
	 * @param flattenTree
	 * @return
	 */
	private static List<ASTCondEQ> extractInnerWhereEQs(
			List<ASTCondEQ> flattenedWhere, List<String> relVars)
	{
		List<ASTCondEQ> retval = new ArrayList<ASTCondEQ>();

		for (ASTCondEQ working : flattenedWhere)
		{
			ASTXQueryVar firstChild = (ASTXQueryVar) working.jjtGetChild(0);
			String var1 = firstChild.varName;
			ASTXQueryVar secondChild = null;

			// If select, just check if firstChild is a relevant variable
			if (working.jjtGetChild(1).id == JJTXQUERYSTRCONST)
			{
				if (relVars.contains(var1))
					retval.add(working);
			}

			// If join, add if both joined variables are relevant
			else
			{
				secondChild = (ASTXQueryVar) working.jjtGetChild(1);
				String var2 = secondChild.varName;
				if (relVars.contains(var1) && relVars.contains(var2))
					retval.add(working);
			}
		}
		return retval;
	}

	/**
	 * Takes two lists of relevant variables (one for each join-for clause), and
	 * an original where clause to analyze, and returns a pair containing the
	 * two join lists
	 * 
	 * @param relVars1
	 *            The relevant vars for the first relation
	 * @param relVars2
	 *            The relevant vars for the second relation
	 * @param inWhere
	 *            The where statement containing the join parameters
	 * @return The list of attribute to join the relations on
	 */
	@SuppressWarnings("unchecked")
	private static List<String>[] findJoinVars(List<String> relVars1,
			List<String> relVars2, ASTWhere inWhere)
	{
		List<ASTCondEQ> flattenedEQs = getEqFromWhere(inWhere,
				new ArrayList<ASTCondEQ>());

		List<String> joinList1 = new ArrayList<String>();
		List<String> joinList2 = new ArrayList<String>();
		for (ASTCond working : flattenedEQs)
		{
			// Only check joins, not selects
			if (working.jjtGetChild(1).id == JJTXQUERYVAR)
			{

				// Get your hands on the variable
				ASTXQueryVar firstChild = (ASTXQueryVar) working.jjtGetChild(0);
				ASTXQueryVar secondChild = (ASTXQueryVar) working
						.jjtGetChild(1);
				String var1 = firstChild.varName;
				String var2 = secondChild.varName;

				if (relVars1.contains(var1) && relVars2.contains(var2))
				{
					joinList1.add(var1.substring(1));
					joinList2.add(var2.substring(1));
				}
				// Equality is commutative, check the other order
				else if (relVars1.contains(var2) && relVars2.contains(var1))
				{
					joinList1.add(var2.substring(1));
					joinList2.add(var1.substring(1));
				}
			}
		}

		List<?>[] retval = new List<?>[] { joinList1, joinList2 };
		return (List<String>[]) retval;
	}

	/**
	 * This helper function flattens a where tree that contains only And and Eq
	 * condition, and returns only the EQ conditions. It recursively returns a
	 * flat list of the tree in in-order traversal order (which should be the
	 * expected ordering of the parse tree)
	 * 
	 * @param working
	 *            The current node in the sbutree to process
	 * @param flattenedTree
	 *            A flat list of EQ statements from the where tree
	 * @return
	 */
	private static List<ASTCondEQ> getEqFromWhere(MySimpleNode working,
			List<ASTCondEQ> flattenedTree)
	{
		if (working != null && working.id == JJTWHERE)
			return getEqFromWhere(working.jjtGetChild(0), flattenedTree);

		// Don't process nodes that aren't either And or Eq nodes
		if (working == null
				|| !(working.id == JJTCONDAND || working.id == JJTCONDEQ))
			return flattenedTree;

		// Traverse the tree
		getEqFromWhere(working.jjtGetChild(0), flattenedTree);
		if (working.id == JJTCONDEQ)
			flattenedTree.add((ASTCondEQ) working);
		getEqFromWhere(working.jjtGetChild(1), flattenedTree);
		return flattenedTree;
	}

	/**
	 * This helper function takes a return clause from the original FLWR and
	 * returns a list of all variables contained in it using breadth first
	 * ordering to be processed using fixJoinVar so that the new return
	 * statement works on the tuple returned by the join
	 * 
	 * It is also used when making the magic list to find out what variables are
	 * in the dependency list of another.
	 * 
	 * @param working
	 *            The node currently being processed recursively
	 * @param varList
	 *            The list of variables being built
	 * @return
	 */
	private static List<ASTXQueryVar> extractVars(MySimpleNode working,
			List<ASTXQueryVar> varList)
	{
		// Stop processing at leaves
		if (working == null)
			return varList;

		if (working.id == JJTXQUERYVAR)
			varList.add((ASTXQueryVar) working);

		for (MySimpleNode child : working.jjtGetChildren())
			extractVars(child, varList);

		return varList;
	}

	/**
	 * Takes in a node in a parse tree, and returns the string representation of
	 * all variables that in the subtree rooted at that node.
	 * 
	 * @param root
	 *            The root of the subtree to check
	 * @return The names of all variables in the subtree
	 */
	private static List<String> extractVarStrings(MySimpleNode root)
	{
		List<String> retval = new ArrayList<String>();

		List<ASTXQueryVar> varNodes = extractVars(root,
				new ArrayList<ASTXQueryVar>());

		for (ASTXQueryVar node : varNodes)
			retval.add(node.varName);

		return retval;

	}

	/**
	 * Takes the return statement of a FLWR statement and rewrites it to a form
	 * that works with the tuple join format
	 * 
	 * @param inReturn
	 *            The outer return to be rewritten
	 */
	private static void makeOuterJoinReturn(ASTReturn inReturn)
	{
		List<ASTXQueryVar> oldVarList = extractVars(inReturn,
				new ArrayList<ASTXQueryVar>());
		List<ASTXQuerySlashRP> newVarList = new ArrayList<ASTXQuerySlashRP>();

		// Create the variable replacements
		for (ASTXQueryVar varNode : oldVarList)
			newVarList.add(fixJoinVar(varNode));

		for (int i = 0; i < oldVarList.size(); i++)
		{
			ASTXQueryVar oldVar = oldVarList.get(i);
			ASTXQuerySlashRP newVar = newVarList.get(i);
			MySimpleNode oldVarParent = oldVar.jjtGetParent();
			int parType = oldVarParent.id;
			boolean onlyVar = (parType != JJTXQUERYSLASHRP && parType != JJTXQUERYDSLASHRP);

			// If var wasn't followed by a path, replace it as parent's child
			if (onlyVar)
			{
				int varChildNum = jjtGetChildNum(oldVarParent, oldVar);
				jjtSetChild(oldVarParent, newVar, varChildNum);
			}

			// Otherwise, replace it's parent as grandparent's child
			else
			{
				MySimpleNode oldVarGParent = oldVarParent.jjtGetParent();
				int parentChildNum = jjtGetChildNum(oldVarGParent, oldVarParent);
				jjtSetChild(oldVarGParent, newVar, parentChildNum);
			}
		}

	}

	/**
	 * This function creates the magic lists by going through the variables, and
	 * creating one list per absolute path. Then it adds a variable to a list if
	 * all variables that it depends on are in that list.
	 * 
	 * @param inFor
	 * @return
	 */
	private static List<List<String>> makeMagicLists(ASTFor inFor)
	{
		List<List<String>> retval = new ArrayList<List<String>>();
		List<ASTXQuery> children = inFor.jjtGetChildren();
		int cIndex = 0;
		int numLists = 0;
		for (; cIndex < children.size(); cIndex++)
		{
			String workingVar = inFor.vars.get(cIndex);
			ASTXQuery working = children.get(cIndex);

			// If absolute path, make a new list for it
			if (working.id == JJTXQUERYAP)
			{
				retval.add(new ArrayList<String>());
				retval.get(numLists).add(workingVar);
				numLists++;
			}
			// Otherwise find out which list it goes into
			else
			{
				List<String> requiredVars = extractVarStrings(working);
				for (List<String> magicList : retval)
				{
					if (magicList.containsAll(requiredVars))
						magicList.add(workingVar);
				}
			}
		}
		return retval;

	}
}
